AI 训练加速原理解析与工程实践分享

本文整理自同名线上分享，是 12 月份「百度百舸 - 云原生 AI」技术公开课的第二期。

这次分享将系统性的分析在 AI 模型训练过程中的主要性能瓶颈，以及当前针对这些瓶颈的主要的加速方案和技术原理，并介绍百度智能云在这方面的一些实践成果。

本系列一共有四期，第三期活动大家可以点击下方按钮预约。

---

今天的分享，主要包括三个部分：‍

首先介绍我们为什么需要做 AI 训练加速，也就是整体背景和出发点是什么；

第二部分我们会系统性的分析实际训练过程中的可能会遇到的性能瓶颈问题，然后针对这些问题，介绍目前主要的加速方案；

第三部分介绍百度百舸平台的 AI 训练加速套件 AIAK-Training 在一些模型训练加速上的实践效果。

1. 为什么需要 AI 训练加速？

在 AI 系统中，一个模型从生产到应用，一般包括离线训练和推理部署两大阶段。

离线训练阶段，就是产生模型的过程，用户需要根据自己的任务场景，准备好训练模型所需要的数据集，以及神经网络算法。

算法可以理解为是一个高度复杂的非凸数学函数，函数中包括很多变量以及参数。模型训练的过程其实就是在学习神经网络模型中的参数。

模型训练开始后，会读取数据，然后送入模型进行前向计算，并计算与真实值的误差。然后执行反向计算得到参数梯度，最后更新参数。训练会进行多轮的数据迭代。

训练完成之后，我们会保存训练好的模型，然后将模型做上线部署，接受用户的真实输入，通过前向计算，完成推理。

因此，无论是训练还是推理，核心都是数据计算。 为了加速计算效率，一般都是通过 GPU 等异构加速芯片来进行训练和推理。

另外，从深度学习模型发展历程来看，为了能够持续突破模型的精度上限，模型参数量其实在快速的膨胀。然而更大的参数量，就会带来更大的计算复杂度。‍‍

下图左侧是摘自一篇公开的论文，从这篇总结里，我们看到在 2010 年之前，模型的计算量大约 20 个月翻一番。在 2010~2015 年，常规模型计算每 5-6 个月翻一番。而在 2015 年之后，衍生了大模型训练的趋势，计算量增长 10~100 倍。

模型训练对算力以及基础设施的要求越来越高，训练需要更多的算力，也需要更长的时间，这也导致了需要更多的资源成本。这里我们列举了一些论文或研究中公开的成本数据，反应了模型训练的费用是非常高昂的。

因此，如何稳定的进行模型训练，如何持续降本增效其实至关重要。

在这样的大背景下，百度智能云推出了百度百舸 · AI 异构计算平台，目标是为 AI 场景提供软硬一体化的解决方案。通过 AI 计算、AI 存储、AI 加速、AI 容器四层技术栈，满足上层业务场景的需求。

• AI 计算层，提供了包括高性能的 GPU、以及昆仑等异构芯片资源，以及高性能的 RDMA 或 IB 网络，以及自研的超级 AI 计算机 X-MAN 等；‍

• AI 存储层，包括对象存储 BOS 满足数据湖存储的需求、以及专为 AI 设计的高性能并行文件系统 PFS；‍

• AI 加速层，包括数据湖存储加速套件 RapidFS，AI 训练加速套件 AIAK-Training，AI 推理加速套件 AIAK-Inference；‍

• AI 容器层，也即是资源调度层，利用云原生的技术能力，满足 GPU、AI 作业等弹性调度的需求。云原生 AI 的内容在我们上一期的技术公开课有专门分享。‍‍

当我们考虑做性能加速的时候，第一个想到的可能是使用更好的硬件。

这会带来一定的性能提升，但是大部分情况下可能并没有充分发挥出硬件的计算能力，核心的原因就是训练代码的执行效率并没有调到最优或更优的状态。

• 首先，框架为了兼容更多的需求，一般会提供一些较为通用的优化能力，更多的会关注易用性、生态建设等；
• 其次，算法工程师在设计算法时，核心的精力在如何提高模型的精度，性能效率关注不足，某些情况下框架提供的一些通用优化能力，也都没有利用起来。

因此，当我们决定使用某种模型算法时，为了达到较好的资源效率和训练效率，我们需要有意识的去优化。不过这里也有很多的技术挑战：

• 性能影响因素比较多，需要根据不同模型自身的运行特点进行分析，没有完全固定的优化方案；

• 性能优化需要理解工程实现的原理，比如当我们去做异构芯片计算的优化，需要专业的异构研发经验才能开展，技术门槛较高；

• 还有些情况，当我们做分布式训练时，可能还需要结合集群硬件和网络拓扑，来优化分布式训练场景下的扩展性问题。

这些挑战极大地影响了模型训练性能的调优，因此我们推出了AIAK-Training 的加速套件，期望通过抽象易用性的接口降低优化成本，并通过软硬协同的优化手段，来充分加速客户在百度智能云上的模型训练性能。

2. 训练性能开销分析和加速方案

在介绍 AIAK-Training 具体效果之前，我们先介绍下训练加速这个话题下关键的技术思路和方案原理是什么样的。

因为模型训练优化本身是一个软硬件综合的工作，技术栈相对比较复杂，今天的内容肯定没办法涵盖全部细节，我们尽量把关键思路讲到。

首先我们看下当前的模型训练方案。过去的发展阶段里，模型训练方案关键有两个层次的变化，一是从单卡训练到分布式训练的变化，二是从数据并行训练到多维混合并行训练的变化。这里的核心驱动点一个是训练的数据量，一个是模型的参数量。

• 单卡训练方式：实际采用这种模式，一般都是模型参数量和数据量相对比较少，单卡的训练时间可以接受。模型参数规模需要保证在训练过程中，单张卡的显存能够满足存储的上限。按照新的 GPU 卡的显存容量配置，一般可以放置的最大规模在10 亿级参数量；
• 当数据集规模比较大的时候，因为训练过程中需要多次遍历全量的数据集，单卡训练就会需要比较长的时间，这时可以从单卡扩展到多卡，通过分布式的方式来加快训练。

这里应用最广泛的就是数据并行。在数据并行方案下，数据集会被平均切分成多份，然后每张卡上保存完整的模型，各自独立并行处理切分后的子数据集。

当模型参数量足够大的时候，比如参数量达到百亿、千亿级别，单卡放不下完整的模型，这里又出现了模型并行、或者同时使用数据并行和模型并行的混合并行方案。

模型并行，会将模型切分到不同的卡上，每个卡上放置模型的一部分，这里又根据切分的方式不同，比如层内切分或层间切分，又细分了 Tensor 并行和流水线并行的方式。‍

因为在一般模型训练中，使用更多的还是数据并行，我们下面重点还是以数据并行为例，来介绍性能优化的思路。

我们从软硬件整体视角先理解下单卡训练过程中存在的性能开销。

下图左边是我们从软件角度上的训练流程。单卡训练的过程，主要包括数据读取、数据预处理，前向计算输出并计算 loss，根据 loss 函数反向计算，得到每一层参数的梯度，最后根据梯度更新模型参数。持续该过程，直到训练收敛。

下图右边，是一个简化的节点硬件拓扑图。最上面是数据存储，可以是本地存储，也可以是网络存储。然后是 CPU、内存，CPU 下通过多个 PCIe Switch 连接着 8 张 GPU 卡，编号从 0~7，8 张卡之间通过 NVSwitch 互联。不同的计算实例，硬件的拓扑结构会有不同。

• 当训练启动时，首先数据读取，涉及从存储介质中读数据到内存中，主要是存储I/O 开销。根据存储介质，要读取的数据量不同，这部分时间开销也不同；
• 接下来是数据预处理，主要是对读入的数据进行一些数据增强的操作，比如对图片进行 resize，对比度、饱和度调整等等，这部分工作大多数情况下是在 CPU 上。当在 CPU 上完成数据预处理操作之后，需要从主机内存拷贝到 GPU 显存上，涉及到主机和设备之间的内存拷贝开销；
• 然后开始前向计算、反向计算、参数更新，这部分更多操作主要是通过 GPU 来进行，主要时间花费在 GPU 计算上面。这个过程里也可能会穿插一些 CPU 上的操作，可能也需要做主机和设备内存之间的拷贝。

因此，从单卡角度看，主要存在 I/O、CPU 预处理、CPU 和 GPU 之间数据拷贝，GPU 计算等方面的开销。

接着我们看下数据并行的过程。

下图左边还是训练的主流程，右边展示了一个 3 机 24 卡的训练集群的硬件拓扑，3 台机器通过网络互联。

前面的部分我们也介绍到了，在数据并行里每个设备并行独立地执行前向和反向计算过程，因此，每个训练进程也都会遇到前面讲的单卡训练中的性能开销问题。

数据并行为了保证和单卡训练在数学上等价，需要确保每张卡的模型参数在迭代过程中始终保持一致。这里一方面需要让各 GPU 卡的模型参数初始化状态一致，这个一般是在训练开始前，通过广播的方式将第一张卡上的参数状态广播到其他的卡。

而在训练期间，由于每个设备处理的数据不同，前向计算所得到的模型损失值也是不同的，因此还需要在每个设备反向计算出梯度之后，进行梯度的平均，以平均后的梯度值来更新模型参数，从而保证每张卡的模型参数在迭代过程中始终保持一致。

梯度平均涉及到通信的过程，包括节点内部卡之间的通信，以及跨节点的网络通信开销。这里的通信又包括同步通信和异步通信，不过为了保证模型训练的收敛，一般都是采用同步通信的方案，后续的优化工作也都是基于同步通信的方式来展开。

由上可知，数据并行相比单卡训练，主要增加了额外的通信开销。

通过前述分析，我们知道加速AI 训练不单是某一方面的工作，需要从数据加载、模型计算、分布式通信等系统维度综合考虑。这里说的数据加载，包括数据 I/O、预处理、内存拷贝等过程。

在具体的优化实践中，给定一个待优化的模型，加速模型训练就是要不断提升训练的总体吞吐（每秒可以训练的样本数）。这个过程，我们一般可以先分析单卡的训练吞吐，当单卡的训练吞吐提升上来之后，我们再看从单卡扩展到多卡，看看如何提升多卡的训练加速比。

首先单卡的训练优化，极限优化目标是全部时间都在 GPU 计算上，加速器利用率100%。当然实际很难完全达到这个状态，但是我们可以按这个指标来牵引或衡量我们的工作。

单卡性能优化关键包括两部分：

• 首先数据加载效率的优化。从存储系统上，我们可以使用更高性能的存储介质，或者基于这些高速存储介质组成的并行文件系统，或者说一些缓存加速系统。前面介绍到的，百度百舸也提供了相应的存储系统方案，比如 PFS、RapidFS 等。除此之外，还需要在框架 dataloader 中优化数据的读取过程。

• 其次就是模型计算效率的优化。主要考虑如何优化计算实现，怎么提升计算单元的利用效率，这块可能就需要结合模型具体分析。

然后从单卡扩展到多卡，目标是如何达到线性加速比。线性加速比这个指标，简单来说就是从 1 张卡扩到 2 张卡训练时，训练的性能是否是单卡的2倍。

这里核心就是优化分布式的通信效率，一方面是硬件层面的优化，另外一方面在实际通信中，需要考虑怎么利用好网络的带宽资源，或者是否能够将通信过程进行隐藏等。

下面我们分别从这几个方面详细展开。

首先是数据加载方面的优化。

当我们实例化一个 dataloader 之后，我们会持续迭代 dataloader 来读取一个 batch 的数据进行模型训练。

如果这里不做任何优化，如下图上半部分所示，每个 batch 的数据加载过程和每个 batch 的模型训练过程，实际上是串行进行的。从 GPU 视角来看，就会出现因为数据加载导致的计算间隙，计算资源存在时间上的浪费。

除了前面讲到我们用更好的硬件直接提升数据读的效率之外，还可以怎么优化呢？

实际在 AI 训练过程中，数据访问上有两个关键特征：

• 当数据集做完 shuffle 之后，每轮训练所需要的 batch 数据以及访问顺序是已知的；
• 任意两个 batch 的数据读可以并行，因为数据之间没有任何依赖关系。

因此，我们在不依赖硬件层面改动的时候，可以做的一个优化工作就是数据预取，当训练第一个 batch 数据的时候，可以提前加载下一个 batch 数据，让 I/O 的过程和 GPU 上的计算充分并行起来。

首先，我们需要利用好 dataloader 中已有的优化方案，一是合理设置 num_workers 超参数，通过多进程的方式读数据，这一步可以实现数据从存储系统中预取到主机内存。二是从主机内存拷贝到 GPU显存，可以通过 pinned memory 的机制来加速。

大概介绍下 pinned memory 加速的主要原理：内存数据有两种类型 pageable memory 和 pinned memory，pageable memory 中的数据有被换出到磁盘上的可能。这种情况下，当执行 H2D 时，可能需要先从磁盘读到内存，然后从内存拷贝到显存。另外，pageable memory 数据拷贝到 GPU 显存时，需要先创建一个临时的 pinned memory 缓冲区，把数据从 pageable memory 拷贝 pinned memory，之后才能传输到 GPU 上，也多了额外的数据搬运的操作。

不过当我们使能了上述方案后，我们仅实现了从存储系统到主机内存的预取，加快了主机到设备的数据拷贝速度。但是主机到设备的内存拷贝，和实际计算 kernel 在 GPU 上还是串行执行，也就是 GPU 上依然存在少量的时间间隙。‍

AIAK 针对这个问题，又做了进一步的优化，可以实现 H2D 和前向计算的 overlap。

在数据并行场景下，还需要注意的一个事情，数据需要均衡的切分。

如果每个训练进程分配的数据不均衡，计算量就会不同，也就导致每个进程前向计算和反向计算完成的时间不同，那么先完成计算的进程，在反向过程中就会先进入到梯度通信环节中，但是因为 Allreduce 通信是同步通信操作，需要所有进程同时开始并同时结束，因此先开始通信的进程，会一直等待其他所有进程也发起了 AllReduce 后才能一起完成通信操作。这里就会出现因为快慢不一导致的资源空闲问题。

为了解决这种问题，需要每个进程使用相同的 batchsize 来读取数据，另外每个 batch 的数据量要均衡。图像类数据一般会固定尺寸进行训练，而像 NLP 类模型需要处理变长的语句，可能需要进行特殊的处理，比如可以将数据 padding 到相同的长度，或者通过样本长度排序的方式来均衡分配等。

下面介绍计算效率的优化。

计算包括前向、反向、参数更新。优化计算的目标，是为了能够充分发挥出异构硬件的算力，理想情况就是让 GPU 芯片实际计算时的性能达到理论峰值。

我们先从一个单算子的角度分析，当我们准备在 GPU 上执行一个计算操作的时候，简化的流程有四步。

• 首先在 CPU 上异步发射一个 GPU 计算 Kernel；

• 当 Kernel 调度执行时，需要先从 GPU 上的 Global Memory 读取计算所需要的数据;

• 开始执行计算;

• 当计算完成后，需要将计算的结果写回 Global Memory。

根据计算、访存的开销占比，一般会将算子分类为计算瓶颈或者访存瓶颈。

当从一个算子扩展到一个完整的模型训练时，因为要连续执行非常多的计算 Kernel，那么 Kernel 计算之间就会出现很多因为 Kernel Launch、中间结果的读写导致的计算间隙问题。

由上可知，优化模型计算效率，需要从访存优化、计算优化、其他开销的优化综合考虑。

• 访存优化，主要考虑如何减少数据在显存和计算单元之间搬运的时间。从算子实现角度上，需要利用好 GPU 存储层次架构，比如把数据搬运到更快的存储器上比如 share memory，减少对 global memory 的访问，从而节省访存时间。或者做好计算指令和访存指令的 overlap，来提升计算访存比。而从单算子扩展到多算子时，还需要考虑如何减少中间结果的读写，这种问题一般可以通过算子融合的手段来优化；

• 计算优化，计算瓶颈问题多半应该是没有正确对任务进行分块，或者没有利用好 GPU 并行计算的优势，导致并行度不高。还有可能没有使用合并指令集导致计算效率低下，或者没有使用 Tensor Core 等高性能计算单元，造成资源浪费。不同的问题，也对应着不同的优化手段；

• Kernel Launch 等其他开销，大量时间花费在访存或计算之外，可以考虑的优化手段如算子融合、Cuda Graph 等。

首先是算子融合。算子在底层 GPU 执行时，会发起一次或者多次 Kernel Launch，Kernel 之间交互数据也需要经过显存，而算子融合就是将多个 GPU Kernel 融合成一个大 Kernel，统一发起和执行。

• 因为减少了需要执行的算子数量，从而可以减少 Kernel 调度和发起的开销；

• 通过融合，可以通过寄存器等来传递中间结果，避免从 global memory 的来回搬运，极大降低了显存等待的时间；

• 在某些场景中，可以通过算子融合，可以更充分的利用计算资源，提升资源和计算的效率。

算子融合具体如何实现呢？

一种方式，分析模型中的低效操作，专家经验手写融合算子。在 GPU 上主要就是 CUDA 算子研发，这里存在一定的门槛。AIAK-Training 会针对典型的模型结构，或者客户需求，提供高效优化的算子实现。

另一种方式，就是编译优化的方案。通过编译的方式进行计算优化，以及代码自动生成，从而降低在不同硬件上的手工优化成本。不过当前很多编译方案更多还是针对推理优化，训练上的方案还在快速演进过程中。不过从极致性能角度看，未来一段时间依然离不开手写融合算子的工作。

下面介绍一些算子融合的实际案例。第一个是针对典型模型网络结构的优化。

下图展示了我们在 SwinTransformer 模型中针对核心模块 WindowAttention 进行的计算融合优化。

WindowAttention 结构，核心操作公式如下图所示。计算过程中，需要依次执行 7 个计算 Kernel。再加上一些 reshape 等转换操作，总共需要 launch 10 个 Kernel。通过性能分析发现，实际执行过程中 launch kernel 的间隔冗余开销占到了端到端 80% 以上的时间，导致该模块存在着较大的优化空间。

通过将这些 Kernel 融合成一个，整个模块的执行时间从 392 微秒减少到 13 微秒，单算子加速了 30 倍。整个模型的训练效率，端到端加速了 20% 以上。

这个优化的核心思路，主要有三点：

• 是利用好 GPU 的三级访存流水：显存、share memory、寄存器；

• 通过分块策略，将 2 个矩阵乘法和 softmax 进行融合；

• 针对前向中间结果的优化，充分利用 Tensor Core，在反向计算中使用重计算代替重加载，使得访存开销极大降低。

下图是一个数据操作的融合举例，是在 FCOS3D 模型中对于坐标压缩操作的一个优化。

通过性能分析发现，这个操作过程中存在大量 GPU 空隙，GPU 利用率较低。该操作主要的功能是根据 index 将 3D-Tensor 压缩成 2D-Tensor。在原生实现中会先在host 端生成 index，再进行 H2D 拷贝，最后完成 Tensor 压缩，这会造成额外的拷贝和等待开销。

为此，我们重新实现这部分操作，核心思路主要就是将操作全部迁移至 GPU 上，直接在 GPU 上完成 index 的生成和 Tensor 的压缩，减少 CPU 参与，同时避免了非必要的 CPU-GPU 之间的内存拷贝。

单算子执行时间从 9.69 毫秒减少到 32 微秒，加速了 300 倍，整个模型端到端训练提升了 10% 以上。

下面我们介绍另一个计算优化的思路，就是提高计算的并行度，充分利用 GPU 并行计算的优势，同样借助一些实际案例来介绍。

我们发现在一些模型中，有一些操作是串行执行的。比如在一些目标检测模型里，在 loss 计算过程中，有些操作并不是按照一个 batch 进行操作，而是 for-loop 每张图片或一个样本，这种情况下，当我们去提升 batchsize 的时候，因为这里的串行，可能没法达到我们想要的性能效果。

以 YOLOv7 中的 SimOTA 操作举例，原生实现中，是通过 for-loop 遍历一个 batch 的每一张图片，然后为图片的 gtbox 执行 SimOTA 标签分配。这种串行的实现方式导致该部分操作的 GPU 利用率非常低效。

而每张图片处理时数据之间是没有依赖的。因此，我们做的一项工作就是将串行计算改成 batch 并行计算，通过对一个 batch 的数据进行并行化的标签分配，从而加速这部分计算的效率。

最终效果上，SimOTA 操作的耗时从 384 毫秒下降到 69 毫秒，计算效率提升 5.5 倍，整个模型的端到端训练效率提升了 18% 以上。

模型训练过程中也有其他的类似场景，比如说参数更新。参数更新时，默认也是通过循环的方式，遍历每个参数，然后每个参数都会启动一个参数更新的 Cuda Kernel，然后依次执行。

针对这种情况，AIAK 也增加了 FusedOptimizer 的优化，通过融合参数更新的算子，实现批量化更新参数，大幅减少 Kernel Launch 次数。

下面介绍了另一个优化手段 CUDA Graph，主要是为了减少 CPU Launch Kernel 的开销。

CUDA Graph 是在 CUDA 10 版本中引入的特性，可以将一系列 CUDA Kernel 封装成单个单元，可以通过一次 CPU Launch 操作来启动多个 GPU Kernel，从而减少了CPU Launch Kernel 的开销。

如下图所示，默认情况下 CPU 需要依次发射多个 Kernel，如果 Kernel 计算时间比较短，Kernel 之间的 Launch 间隙可能成为性能瓶颈。通过 CUDA Graph，仅需要花费一些额外时间构建 Graph，后续通过一次 Graph 的发射，即可大幅缩短实际执行时 Kernel 之间的间隙。

现在很多框架也都增加了对 CUDA Graph 的支持，通过插入一些代码来使K能这个功能。不过也有一些使用上的限制，比如不支持动态 shape、不支持动态控制流、过程中不能捕获 CPU 操作等等，可以根据模型情况，尝试使用下这种优化能力。

下面介绍最后一个计算优化手段，充分利用 Tensor Core 计算单元。

一个 GPU 内一般包含多个 SM，每个 SM 包括不同数据类型的计算核心，以及各类存储资源等。下图左边所示，是一个 NVIDIA A100 SM 的示意图，一个 SM 包含 64 个 FP32 CUDA Core，以及 4 个 Tensor Core。

模型训练时，默认情况主要是用到 FP32 CUDA Core 计算，而 Tensor Core 是一块特殊的硬件执行单元，是从 Volta 系列 GPU 开始引入，主要是用来加速矩阵或卷积的操作效率。

相比 FP32 CUDA Core 一次只能执行两个标量的计算，Tensor Core 可以一次对两个矩阵执行计算，因此 Tensor Core 的计算吞吐远高于 FP32 CUDA Core。

在 A100 中，Tensor Core 支持了多种浮点数据类型，对于深度学习训练来说，可能涉及到包括 FP16、BF16、以及 TF32 模式。

TF32 主要用于单精度训练场景，相比 FP32 训练，在保持相同的访存带宽需求下，理论计算吞吐量提高了 8 倍。

FP16/BF16 主要用于混合精度训练场景，相比 FP32 训练，访存需求减少了一半，理论计算吞吐量提高了 16 倍。

使用 Tensor Core 的方式，可以使用底层的 cublas 或 cuda 接口进行编程，而对于算法开发者来说，更直接的就是使用框架中提供的 TF32 训练或混合精度训练方案。

首先是 TF32 训练模式，TF32 是 Ampere 开始引入。

TF32 在浮点数的表达中，有 8 个指数位，10 个尾数位和 1 个符号位。指数位与  FP32 相同，即数据表示范围相同，但是尾数位低于 FP32，和 FP16 相同。

需要注意的是，TF32 不是一个对外开放的数值类型，只是 Tensor Core 的一种计算模式，也就是用户不能去直接创建一个 TF32 类型的浮点数。

当使能 TF32 的时候，Tensor Core 计算矩阵或卷积操作时，会自动将 FP32 转换成 TF32，计算完成之后，输出的数据类型依然是 FP32 类型。

TF32 训练在某些框架版本中是默认开启，某些框架版本中可能需要通过环境变量或者参数配置来手工开启，具体需要参考框架的用户手册。

不过由于 TF32 相比 FP32 来说，精度范围降低了，实际训练时还需要关注对模型收敛精度的影响。

混合精度训练是指在尽可能减少模型精度损失的情况下，使用 FP32 和 FP16 混合精度进行训练。

混合精度训练的收益主要有：相比 FP32 训练，内存需求减少，可以训练更大的网络或使用更大的 batchsize。使用更少的内存带宽，可以加速数据传输，半精度的计算也可以让数学运算效率更快；

不过因为 FP16 的指数位和尾数位的范围都比 FP32 要少，因此数值表示范围和精度都会有降低，在实际使用的时候，就可能出现因为表示范围狭窄导致的数值溢出问题，或者因为精度不足导致舍入误差。

为了优化类似问题，混合精度训练方案中有几个关键的技术工作：

• 算子黑白名单机制，框架使用黑白名单自动为算子选择精度，模型训练过程中，会自动的插入 cast 操作进行类型转换，不需要开发者干预。针对数值精度敏感的计算，依然使用 FP32 来算，而对于数值安全的计算，比如矩阵乘，则会使用  FP16 来计算；

• 训练过程中，会存储一份 FP32 的权重参数，用于训练时候的参数更新，优化舍入误差的问题；

• 针对 FP16 容易溢出的问题，使用了 Loss scaling 的方案，比如对 Loss 放大 n 倍，根据链式法则，梯度也会随之放大 n 倍，从而使其落到 FP16 的表示范围内，具体过程如下图左侧所示。

目前所有框架都已经支持混合精度。AIAK-Training 组件进一步引入 NVIDIA Apex 中  AMP O2 混合精度模式，这个模式会更加激进的将更多计算转 FP16 来加速训练。相比默认 O1 模式，速度会有进一步提升，不过精度可能会受影响，需要结合具体模型验证。

AIAK-Training 提供兼容 torch amp 原生用法的使用方式，方便使能 O2 模式。

下面介绍通信优化，这个也是一个非常大的话题，涉及到内容也非常多。

前面也介绍了，通信主要是分布式训练中引入的，因为从单卡扩展到多卡，需要进行多卡之间的一些数据同步，这些同步操作就是通过通信来实施的。

下图列了一个通信优化的整体架构：‍

• 最底层就是网络协议层，包括传统的 TCP 网络、以及在训练场景用的越来越多的 RoCE 或 IB 的 高性能 RDMA 网络。这些底层网络方案，百度百舸也都提供了支持。显然通过改善硬件基础设施来提升网络带宽、降低延迟，是最直接有效的优化通信性能的方法。

• 然后是通信库层。因为需要使用底层的网络协议来进行通信，并且实际应用时可能涉及多种通信原语，比如点对点通信、集合通信等，因此也出现一些高度封装并且优化的通信库。在 GPU 训练场景中，我们一般都是使用 NCCL 通信库，性能比较优。

• 基于底层通信库，上层框架可以比较方便的构建分布式训练的通信架构。常见的包括参数服务器架构，集合通信架构。目前 CV/NLP 等领域的模型，主要都是采用集合通信的架构方式。

• 通信策略层，主要是在应用层上做一些通信效率的优化，比如通信隐藏、通信融合、通信压缩、通信降频等不同的思路，这些方式大部分也可以叠加一起使用。

我们先看通信策略层面的优化思路，首先是通信隐藏优化。

在数据并行中，梯度同步通信是在训练的反向过程中进行的，当反向算出梯度之后，就可以进行全局的梯度平均。

如果不做任何的机制优化，反向计算和通信就会串行的进行，就会存在计算上的时间间隙。

由于反向过程中，上一个梯度的通信和下一个梯度的计算，两者之间没有任何数据依赖，因此可以让上一个梯度通信和下一个梯度计算并行起来，让两者的耗时相互重叠，从而可以隐藏部分的通信耗时。

在实现层面上，通常是将通信和计算算子调度到不同的 cuda 流上实现的，通信算子调度到通信流，计算算子调度到计算流，不同流上的算子可以并行发射执行，从而实现反向中梯度通信和计算的并行重叠。

目前这个优化能力，在框架中都是默认开启的。

其次，通信融合优化。

默认情况下，模型中的每个梯度都需要发起一次通信操作。如果单个梯度的 size 比较小，那么小数据包在实际通信时，网络带宽的利用率就会非常低，通信性能较差。

通信融合，就是将将多个梯度融合到一起统一进行一次通信，从通信开销模型上分析，既能提升带宽利用率，又能减少通信的初始化延迟项。

现在很多分布式训练框架中，也默认都支持梯度融合的策略，不同框架实现方式有一定的区别，有的实现需要先进行梯度协商确定通信顺序，有的则是直接通过静态的通信分桶。

虽然框架中默认支持通信融合，但是梯度融合的大小一般可以通过参数来配置，用户可以根据物理环境和模型的需求，调节合适的融合阈值，应该可以取得更佳的收益。

如果在网络带宽较低的训练场景中，比如低带宽的 TCP 环境中，梯度同步的延迟可能会成为训练的主要性能瓶颈。

• 这种情况下，可以考虑的一个优化手段就是通信压缩。通信压缩，主要有三种不同的压缩思路：

• 量化压缩，比如使用更低精度来表示梯度，这种方式的压缩率较低，最大能从  32 位压缩到 1 位，也就是最大压缩 32 倍。

• 稀疏化压缩，典型的算法比如 DGC 算法，核心思想是每轮迭代只传输重要的梯度，也就是梯度数值超过设定的某一个阈值，同时为了减少信息损失，会把剩下不重要的梯度在本地进行累积，只要时间足够，最终累积梯度就会超过所设定的阈值，再进行梯度交换。通过这种方式减少通信的数据量，降低对网络带宽的需求。同时为了减少对模型收敛的影响，还通过动量校正、梯度裁剪、动量因子掩蔽、预热训练等不同方式来缓解。这种方案目前主要支持 SGD 优化器。

• 低秩矩阵压缩方式，典型的算法比如 PowerSGD，核心思路是将一个大的梯度矩阵分解成多个小梯度矩阵，通过传输小矩阵来减少通信量。

通信降频优化，最简单的思路就是增大 batchsize，每次迭代更多的数据，减少了迭代次数，也就是减少了通信量。

不过 batchsize 也不是越大越好，越大的 batchsize 可能会导致模型收敛精度下降或者收敛速度变慢。针对类似问题，业界也提出比如 LARS、LAMB 优化器算法，通过分层自适应调节学习率，缓解类似问题，AIAK-Training 也增加了支持。

为了增大 batchsize，如果显存比较充足，可以直接调整 batchsize 超参。如果显存比较吃紧，还可以通过梯度累加的方式，跳过若干次梯度通信，实际也相当于增大了 batchsize。

下面介绍了一种针对通信拓扑的优化方案——分层拓扑通信，也是针对机间网络带宽比较低的情况。

通过分层通信，可以充分的利用机内的高的互联带宽，同时弱化机间低网络带宽的影响。

AIAK 中也实现了这种通信方案，在 25Gbps TCP 环境下，实测 4 机 32 卡 SwinTransformer 训练，通过分层 allreduce，性能可以加速 85%。

最后介绍一个底层通信库层面的优化，GPU Direct RDMA 通信技术，这个技术需要硬件环境上支持 RDMA 网络。

RDMA 通信，允许本地应用程序直接读写远程应用程序的用户态虚拟内存，整个通信过程，除了最一开始提交发送请求这一步需要 CPU 的参与，其余都是由网卡硬件完成的，不需要内存拷贝、系统中断以及软件处理，因此可以实现极致的低时延和高带宽。

而在 GPU 场景中，GPU Direct RDMA 技术进一步增加了 RDMA 直接访问 GPU 显存的支持，避免通信的时候，数据在 GPU 显存和主机内存中来回拷贝，跨机通信延迟进一步降低。

不过在实际的案例中，我们发现有些用户购买了 RDMA 的环境，但是实际并没有用起来 GDR 技术，导致通信效率不高。这里列出了几个关键的配置项，如果有类似问题的话，可以依次进行排查和设置。

前面介绍了当前的一些主要的性能优化思路和方案，整体来看，无论 I/O 优化、计算优化、通信优化，最朴素的优化思路主要就是如何优化操作本身、或者是否可以减少操作发生的次数，或者是否可以操作和其他的过程并行起来从而隐藏开销等。

3. AIAK-Tranining 加速套件实践

前面介绍了很多的优化工作，要想正确的使能，需要每个用户对于框架的工程实现原理比较清晰。为了简化训练优化的成本，我们构建了 AIAK-Training 这个加速套件。

AIAK-Training 会围绕数据加载、模型计算、通信等方面，构建全链路优化能力，同时我们会把这种优化能力，封装成简单易用的接口，用户插入几行代码，即可比较方便的集成使用。同时，我们也在构建自动化策略组合调优的机制，自动帮助用户选择有效的优化策略。

具体使用时，加速库组件可以独立的安装部署，也可以直接使用我们提供的容器镜像。

下面是一些具体的应用案例。

下图所示，主要是针对 dataloader 的优化。这个场景中，模型比较小，数据集规模也比较小，纯计算的速度其实比较快，但是跨 EPOCH 的数据加载时间比较长，导致 I/O 耗时成为了主要的瓶颈。

通过使用 AIAK 中提供的进程复用、以及充分预取机制，整个模型训练加速了  166%。

下图是一个针对模型计算优化的案例。

Transformer 类模型训练场景，实际训练时通信扩展性接近线性，I/O 耗时占比也非常低，计算是主要的性能瓶颈。

针对这个模型，AIAK-Training 进行了一系列计算层面的优化，包括主要结构的算子融合、混合精度、大 batch 调优等等，整个模型的训练效率提升了 169%。

下图的案例主要是应用了通信层面的优化，在云上 TCP 环境中使能针对低带宽网络的优化策略，在一些经典模型上比如 resnet50、bert、vgg16，可以加速 26%~78%。

在自动驾驶场景中，我们也针对典型的 2D 视觉、3D 视觉、激光雷达，以及前融合类的模型，做了一系列的模型训练性能优化，训练性能加速 49%~ 391%。如果对这一类算法的加速感兴趣，也欢迎后续和我们交流。

- - - - - - - - - - END - - - - - - - - - - 

点击阅读原文，了解 云原生AI 更多内容

传送门